/*
 * Grammar.java
 *
 * Created on June 26, 2006, 9:54 AM
 *
 */

package galronnlp.util;

import galronnlp.cfgEntry;
import galronnlp.pcfg.grammar.*;
import galronnlp.pcfg.lexicon.*;
import galronnlp.cfg.grammar.malformedGrammarException;
import galronnlp.cfg.lexicon.malformedLexicalEntryException;
import galronnlp.io.*;
import galronnlp.util.Symbol;
import galronnlp.util.SymbolGenerator;
import galronnlp.util.Tree;
import galronnlp.util.BinaryTree;
import galronnlp.util.UnaryTree;
import galronnlp.util.LexicalTree;
import galronnlp.util.NaryTree;
import galronnlp.math.*;

import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Vector;
import java.util.LinkedList;
import java.util.List;
import java.util.Iterator;
import java.lang.Double;

/**
 * This implementation of a grammar data structure for modified Chomsky Normal
 * Form grammars (CNF grammars with only unary, binary, and lexical rules) is
 * designed for a bottom-up probabilistic CKY parser.
 *
 * @author Daniel A. Galron
 */
public class Grammar {
    
    private Symbol StartSymbol = null;
    
    private LinkedList<galronnlp.cfg.grammar.Rule> productions = new LinkedList<galronnlp.cfg.grammar.Rule>();
    private GrammarReader reader;
    private LexiconReader lreader;
    
    public Hashtable<Symbol, LinkedList<Tuple>> unaryRules        = new Hashtable<Symbol, LinkedList<Tuple>>();
    public Hashtable<SymbolTuple, LinkedList<Tuple>>  binaryRules = new Hashtable<SymbolTuple, LinkedList<Tuple>>();
    public Hashtable<Symbol, LinkedList<Tuple>> lexicon           = new Hashtable<Symbol, LinkedList<Tuple>>();
    
    private Matrix unaryClosure = null;
    private Hashtable<Symbol, Integer> nonTerminalIndex = new Hashtable<Symbol, Integer>();
    
    
    public double score(Tree t) {
        LinkedList<Tree> subtrees = t.children;
        double score = 1;
        if(t.type() == Tree.BINARY) {
            score *= score((Tree)t.children.getFirst());
            score *= score((Tree)t.children.getLast());
            BinaryRule bin = new BinaryRule(t.cat(), t.RHS(), 0.0);
            score *= lookup(bin);
        } else if(t.type() == Tree.UNARY) {
            score *= score((Tree)t.children.getFirst());
            UnaryRule un = new UnaryRule(t.cat(), t.RHS(), 0.0);
            score *= lookup(un);
        } else if(t.type() == Tree.LEXICAL) {
            Word lex = new Word(t.cat(), t.RHS().getFirst(), 0.0);
            score *= lookup(lex);
        } else {
            System.err.println("N-ary not yet supported");
        }
        return score;
    }    
    public String toString() {
        String ret = "lexicon:\n";
        ret += "--------------------------------------------\n";
        for(Enumeration<Symbol> e = lexicon.keys(); e.hasMoreElements(); ) {
            Symbol word = e.nextElement();
            List<Tuple> els = lexicon.get(word);
            for(Iterator<Tuple> i = els.iterator(); i.hasNext(); ) {
                Tuple el = i.next();
                Symbol cat = (Symbol)el.get(0);
                double prob = ((Double)el.get(1)).doubleValue();
                ret += "| " + cat + " :- " + word + ": " + prob + "\n";
            }
        }
        ret += "--------------------------------------------\n";
        ret += "unary rules:\n";
        ret += "--------------------------------------------\n";
        for(Enumeration<Symbol> e = unaryRules.keys(); e.hasMoreElements(); ) {
            Symbol RHS = e.nextElement();
            List<Tuple> els = unaryRules.get(RHS);
            for(Iterator<Tuple> i = els.iterator(); i.hasNext(); ) {
                Tuple el = i.next();
                Symbol LHS = (Symbol)el.get(0);
                double prob = ((Double)el.get(1)).doubleValue();
                ret += "| " + LHS + " --> " + RHS + ": " + prob + "\n";
            }
        }
        ret += "--------------------------------------------\n";
        ret += "binary rules:\n";
        ret += "--------------------------------------------\n";
        for(Enumeration<SymbolTuple> e = binaryRules.keys(); e.hasMoreElements(); ) {
            SymbolTuple RHS = e.nextElement();
            Symbol lRHS = (Symbol)RHS.get(0);
            Symbol rRHS = (Symbol)RHS.get(1);
            List<Tuple> els = binaryRules.get(RHS);
            for(Iterator<Tuple> i = els.iterator(); i.hasNext(); ) {
                Tuple el = i.next();
                Symbol LHS = (Symbol)el.get(0);
                double prob = ((Double)el.get(1)).doubleValue();
                ret += "| " + LHS + " --> " + lRHS + ", " + rRHS + ": " + prob + "\n";
            }
        }
        ret += "--------------------------------------------\n";
        return ret;
    }
    public Grammar(String grammarString, String lexiconString, Symbol StartSymbol, SymbolGenerator gen, boolean string) {
        this.StartSymbol = StartSymbol;
        reader  = new GrammarReader(grammarString, gen, string);
        lreader = new LexiconReader(lexiconString, gen, string);
        this.readInGrammar();
    }
    /** Creates a new instance of Grammar */
    public Grammar(String grammarFileName, String lexiconFileName, Symbol StartSymbol, SymbolGenerator gen) {
        this.StartSymbol = StartSymbol;
        reader  = new GrammarReader(grammarFileName, gen);
        lreader = new LexiconReader(lexiconFileName, gen);
        this.readInGrammar();
    }
    private void readInGrammar() {
        while(reader.hasNext) {
            try {
                Rule rule = (Rule)reader.readNext();
                if(rule != null) {
                    productions.addLast(rule);
                    Tuple element = new Tuple(2);
                    element.put(0, rule.LHS());
                    element.put(1, new Double(rule.prob()));
                    if(rule.arity() == 1) {
                        if(!unaryRules.containsKey(rule.RHS().get(0))) {
                            LinkedList values = new LinkedList();
                            values.addLast(element);
                            unaryRules.put(rule.RHS().get(0), values);
                        } else {
                            ((LinkedList)unaryRules.get(rule.RHS().get(0))).addLast(element);
                        }
                    } else if(rule.arity() == 2) {
                        SymbolTuple key = new SymbolTuple(rule.RHS().get(0), rule.RHS().get(1));
                        if(!binaryRules.containsKey(key)) {
                            LinkedList values = new LinkedList();
                            values.addLast(element);
                            binaryRules.put(key, values);
                        } else {
                            ((LinkedList)binaryRules.get(key)).addLast(element);
                        }
                    } else {
                        System.err.println("Grammar should be in modified CNF");
                        System.exit(1);
                    }
                }
            } catch(malformedGrammarException mge) {
                mge.printStackTrace();
            }
        }
        while(lreader.hasNext) {
            try {
                LexicalEntry lex = (LexicalEntry)lreader.readNext();
                if(lex != null) {
                    productions.addLast(lex);
                    Tuple element = new Tuple(2);
                    element.put(0, lex.category());
                    element.put(1, lex.prob());
                    if(!lexicon.containsKey(lex.entry())) {
                        LinkedList<Tuple> values = new LinkedList<Tuple>();
                        values.add(element);
                        lexicon.put(lex.entry(), values);
                    }
                    else {
                        ((LinkedList)lexicon.get(lex.entry())).addLast(element);
                    }
                }
            } catch(malformedLexicalEntryException mlee) {
                mlee.printStackTrace();
            }
        }
        
    }
    public double lookup(Rule rule) {
        int arity = rule.arity();
        Symbol LHS = rule.LHS();
        List<Symbol> RHS = rule.RHS();
        for(Iterator<galronnlp.cfg.grammar.Rule> i = productions.iterator(); i.hasNext(); ) {
            galronnlp.pcfg.grammar.Rule r = (galronnlp.pcfg.grammar.Rule)i.next();
            if(LHS.equals(r.LHS()) && RHS.equals(r.RHS()) && (arity == r.arity())) {
                return r.prob();
            }
        }
        return 0.0;
    }
    public double lookup(LexicalEntry entry) {
        Symbol LHS = entry.LHS();
        Symbol RHS = (Symbol)entry.RHS().get(0);
        LinkedList<Tuple> lex = lexicon.get(RHS);
        for(Iterator<Tuple> it = lex.iterator(); it.hasNext(); ) {
            Tuple next = it.next();
            Symbol cat = (Symbol)next.get(0);
            if(cat.equals(LHS))
                return ((Double)next.get(1)).doubleValue();
        }
        return -1.0;
    }
    public Symbol startSymbol() {
        return this.StartSymbol;
    }
    public List<Tuple> getLexicalEntry(String word) {
        return lexicon.get(word);
    }
    public static void main(String[] args) {
        SymbolGenerator gen = new SymbolGenerator();
        Grammar gram = new Grammar("grammar.txt", "lexicon.txt", gen.create("S"), gen);
        System.out.println(gram);
    }
    
//    public void calculateUnaryClosure() {
//        for(Enumeration<Symbol> e = unaryRules.keys(); e.hasMoreElements(); ) {
//            
//        }
//    }
//    
//    private 
}
